# server.py

from fastmcp import FastMCP
from db import db_pool

# Create an MCP server
mcp = FastMCP("postgis-mcp")


@mcp.tool()
def get_all_tables() -> list[str]:
    """获取数据库所有表名"""
    with db_pool.get_connection() as conn:
        with conn.cursor() as cursor:
            cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema='public'")
            return [row[0] for row in cursor.fetchall()]
    db_pool.put_connection(conn)

@mcp.tool()
def get_table_columns(table_name: str) -> list[str]:
    """获取表所有列名和结构"""
    with db_pool.get_connection() as conn:
        with conn.cursor() as cursor:
            cursor.execute(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name}'")
            return [row[0] for row in cursor.fetchall()]
    db_pool.put_connection(conn)
        
@mcp.tool()
def get_table_data(table_name: str, columns: list[str], limit: int = 10) -> list[dict]:
    """获取表前limit条数据"""
    with db_pool.get_connection() as conn:
        with conn.cursor() as cursor:
            cursor.execute(f"SELECT {', '.join(columns)} FROM {table_name} LIMIT {limit}")
            return [dict(zip(columns, row)) for row in cursor.fetchall()]
    db_pool.put_connection(conn)
@mcp.tool()
def select_sql(sql: str) -> list[dict]:
    """
    通用select语句
    """
    if not sql.strip().startswith("SELECT"):
        raise ValueError("通用查询语句只支持select语句")
    with db_pool.get_connection() as conn:
        with conn.cursor() as cursor:
            cursor.execute(sql)
            return [dict(zip(cursor.description, row)) for row in cursor.fetchall()]
    db_pool.put_connection(conn)
    
if __name__ == "__main__":
    mcp.run()